import matplotlib
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

matplotlib.rcParams['font.family'] = 'Microsoft YaHei'  # Windows 常用

# -------------------------------------------------
# 1️⃣  原始 3×3 像素块的 RGB 值
#    (左上 → 右下 按行优先排列)
# -------------------------------------------------
c00 = np.array([128,  64, 128])   # (0,0)
c10 = np.array([128,  64, 255])   # (1,0)
c20 = np.array([255,  64, 255])   # (2,0)

c01 = np.array([256,  64, 128])   # (0,1)
c11 = np.array([256,  64, 256])   # (1,1)
c21 = np.array([128,  64, 256])   # (2,1)

c02 = np.array([256, 128, 128])   # (0,2)
c12 = np.array([256, 128, 256])   # (1,2)
c22 = np.array([128, 128, 256])   # (2,2)

# -------------------------------------------------
# 2️⃣  一维 Lagrange 插值（支持任意数量的基点）
# -------------------------------------------------
def lagrange_1d(xs, ys, t):
    """
    xs : 基点的归一化坐标数组（长度 n）
    ys : 对应的数值（标量或向量），形状 (n,) 或 (n,3)
    t  : 目标归一化坐标（标量或数组），返回插值结果
    """
    n = len(xs)
    t = np.asarray(t)
    # 计算 Lagrange 基函数
    L = np.ones((n, ) + t.shape, dtype=float)   # (n, …)
    for i in range(n):
        for j in range(n):
            if i != j:
                L[i] *= (t - xs[j]) / (xs[i] - xs[j])
    # 加权求和
    result = np.tensordot(L, ys, axes=([0], [0]))
    return result

# -------------------------------------------------
# 3️⃣  二维 Lagrange 插值（先 x 再 y）
# -------------------------------------------------
def lagrange_grid_3x3(c00, c10, c20,
                     c01, c11, c21,
                     c02, c12, c22,
                     scale):
    """
    对 3×3 的九个已知像素做二维 Lagrange 插值。
    返回 shape = (scale+1, scale+1, 3) 的 RGB 矩阵。
    """
    # 归一化坐标 0、0.5、1 对应 3 个基点
    xs = np.linspace(0, 1, 3)   # [0., 0.5, 1.]
    ys = np.linspace(0, 1, 3)

    # 把九个点按行组织成 (3,3,3) 的数组
    grid = np.array([
        [c00, c10, c20],
        [c01, c11, c21],
        [c02, c12, c22]
    ])   # shape (3,3,3)

    # 目标插值坐标
    tx = np.linspace(0, 1, scale + 1)
    ty = np.linspace(0, 1, scale + 1)

    # 先在 x 方向对每一行做 1‑D Lagrange 插值，得到 (3, scale+1, 3)
    interp_x = np.empty((3, scale + 1, 3), dtype=float)
    for row in range(3):
        interp_x[row] = lagrange_1d(xs, grid[row], tx)

    # 再在 y 方向对得到的三条曲线做插值，得到最终 (scale+1, scale+1, 3)
    result = np.empty((scale + 1, scale + 1, 3), dtype=float)
    for col in range(scale + 1):
        # 对每个颜色通道分别插值
        result[:, col, :] = lagrange_1d(ys, interp_x[:, col, :], ty)

    return result

# -------------------------------------------------
# 4️⃣  生成 4 倍插值结果（得到 5×5 网格，实际像素是 4×4）
# -------------------------------------------------
scale = 4
interp_rgb = lagrange_grid_3x3(
    c00, c10, c20,
    c01, c11, c21,
    c02, c12, c22,
    scale
)   # shape (5,5,3)

# -------------------------------------------------
# 5️⃣  绘图函数（保持不变，只是标题改为 3×3 示例）
# -------------------------------------------------
def plot_interpolation(original, interp, scale):
    """
    original: 3×3×3 原始颜色矩阵（仅用于演示，可不显示）
    interp:   (scale+1)×(scale+1)×3 插值后颜色矩阵
    """
    fig, ax = plt.subplots(figsize=(6, 6))
    ax.set_aspect('equal')
    ax.set_xlim(-0.5, scale + 0.5)
    ax.set_ylim(-0.5, scale + 0.5)
    ax.invert_yaxis()                     # y 轴向下为正（像素坐标习惯）

    # ---- 绘制放大后 4×4 像素块（每块大小 1×1，坐标从 0 到 4） ----
    for i in range(scale):
        for j in range(scale):
            ij = np.clip(interp[i, j], 0, 255)      # 保证合法
            color = ij / 255.0
            rect = Rectangle((j, i), 1, 1, facecolor=color, edgecolor='none')
            ax.add_patch(rect)

            ax.text(j + 0.5, i + 0.5,
                    f'({int(ij[0])},{int(ij[1])},{int(ij[2])})',
                    ha='center', va='center', fontsize=8, color='white')

    # ---- 网格线 ----
    ax.set_xticks(np.arange(-0.5, scale + 1, 1), minor=True)
    ax.set_yticks(np.arange(-0.5, scale + 1, 1), minor=True)
    ax.grid(which='minor', color='gray', linewidth=0.5, linestyle='--')

    # ---- 标注 ----
    ax.set_title(f'牛顿‑拉格朗日 3×3 插值放大 {scale} 倍示意图', fontsize=14)
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    plt.show()

# -------------------------------------------------
# 6️⃣  组织原始 3×3 矩阵（仅作演示，可不参与绘图）
# -------------------------------------------------
original_rgb = np.stack([
    [c00, c10, c20],
    [c01, c11, c21],
    [c02, c12, c22]
])   # shape (3,3,3)

# -------------------------------------------------
# 7️⃣  绘制结果
# -------------------------------------------------
plot_interpolation(original_rgb, interp_rgb, scale)
